#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/2/26 21:25
# @Author  : Seven
# @File    : PositionMeasurement.py
# @Software: PyCharm
# function : 实现位姿测量算法
import cv2
import numpy as np
import glob

# 加载相机标定的数据
with np.load('C.npz') as X:
    mtx, dist, _, _ = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')]


def draw(img, corners, imgpts):
    """
    在图片上画出三维坐标轴
    :param img: 图片原数据
    :param corners: 图像平面点坐标点
    :param imgpts: 三维点投影到二维图像平面上的坐标
    :return:
    """
    corner = tuple(corners[0].ravel())
    cv2.line(img, corner, tuple(imgpts[0].ravel()), (255, 0, 0), 5)
    cv2.line(img, corner, tuple(imgpts[1].ravel()), (0, 255, 0), 5)
    cv2.line(img, corner, tuple(imgpts[2].ravel()), (0, 0, 255), 5)
    return img


# 初始化目标坐标系的3D点
objp = np.zeros((6 * 7, 3), np.float32)
objp[:, :2] = np.mgrid[0:7, 0:6].T.reshape(-1, 2)
# 初始化三维坐标系
axis = np.float32([[3, 0, 0], [0, 3, 0], [0, 0, -3]]).reshape(-1, 3)  # 坐标轴
# 加载打包所有图片数据
images = glob.glob('image/*.jpg')
for fname in images:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 找到图像平面点坐标点
    ret, corners = cv2.findChessboardCorners(gray, (7, 6), None)
    if ret:
        # PnP计算得出旋转向量和平移向量
        _, rvecs, tvecs, _ = cv2.solvePnPRansac(objp, corners, mtx, dist)
        print("旋转变量", rvecs)
        print("平移变量", tvecs)
        # 计算三维点投影到二维图像平面上的坐标
        imgpts, jac = cv2.projectPoints(axis, rvecs, tvecs, mtx, dist)
        # 把坐标显示图片上
        img = draw(img, corners, imgpts)
        cv2.imshow('img', img)
        cv2.waitKey(500)

cv2.destroyAllWindows()
